import os  # 导入操作系统相关模块
import argparse  # 导入命令行参数解析模块
import pandas as pd  # 导入pandas库，用于数据处理
import numpy as np  # 导入numpy库，用于数值计算
from NARM import pretrain_narm  # 导入NARM中的预训练函数
from time import time  # 导入时间模块，通常用于记录时间
import warnings  # 导入警告模块

# 关闭所有警告
warnings.filterwarnings("ignore")

# 设置命令行参数解析器
parser = argparse.ArgumentParser()
# 添加'--offline'参数，用于指示是否使用离线数据
parser.add_argument('--offline', action='store_true', help='offline')
# 添加'--epochs'参数，用于设置训练的轮数，默认为1
parser.add_argument('--epochs', type=int, default=1, help='the number of epochs to train for')
# 解析命令行参数
opt = parser.parse_args()
# 打印解析的参数
print(opt)

# 加载训练数据的函数
def load_data(train_path):
    if len(train_path) == 2:
        # 如果train_path中有两个文件，分别读取并合并
        train_click1 = pd.read_csv(train_path[0])
        train_click2 = pd.read_csv(train_path[1])
        train_click = pd.concat([train_click1, train_click2], axis=0)
    else:
        # 如果只有一个文件，直接读取
        train_click = pd.read_csv(train_path)

    return train_click

# 主程序入口
if __name__ == '__main__':
    # 根据命令行参数判断是否使用离线数据
    if opt.offline:
        data_path = 'offline_data'  # 离线数据路径
    else:
        data_path = '../tcdata'  # 在线数据路径

    # 构造训练数据的路径
    train_path = ['{}/train_click_log.csv'.format(data_path), '{}/testA_click_log.csv'.format(data_path)]
    # 加载数据
    train_click = load_data(train_path)

    # 使用NARM进行预训练
    pretrain_narm(train_click, opt.epochs)
    print('[-] PRETRAINED OF NARM FINISHED.\n')
